#include "EnhanceWrapper.h"

EnhanceWrapper::EnhanceWrapper(std::string modelPath){
    double tS, tE;

    tS = (double) clock();
    std::cout << "init cann runtime..." << std::endl;
    this->initCannContext(modelPath);

    std::cout << "init model memory" << std::endl;
    this->initModelResource();
    tE = (double) clock();
    std::cout << "init time cost: " << (tE - tS) / CLOCKS_PER_SEC << " s" << std::endl;
}

cv::Mat EnhanceWrapper::infer(std::string imageSrc){
    double tS, tE;
    tS = (double) clock();
    std::pair<cv::Mat, cv::Size> transPair = this->preprocess(imageSrc);
    tE = (double) clock();
    std::cout << "Preprocess Done, cost: " << (tE - tS) / CLOCKS_PER_SEC << " s" << std::endl;

    tS = (double) clock();
    std::vector<float> buffer = this->_infer(transPair.first);
    tE = (double) clock();
    std::cout << "Infer Done, cost: " << (tE - tS) / CLOCKS_PER_SEC << " s" << std::endl;

    tS = (double) clock();
    cv::Mat res = this->postProcess(buffer, transPair.second);
    tE = (double) clock();
    std::cout << "PostProcess Done, cost: " << (tE - tS) / CLOCKS_PER_SEC << " s" << std::endl;

    return res;
}

cv::Mat EnhanceWrapper::postProcess(const std::vector<float> &buffer, cv::Size oriSize){
    std::vector<float> outBuffer(buffer.size());
    NetworkTool::toHWC(buffer, outBuffer, cv::Size(512, 512));

    cv::Mat floatMat(512, 512, CV_32FC3, outBuffer.data());
    floatMat.convertTo(floatMat, CV_8UC3);
    cv::resize(floatMat, floatMat, oriSize, cv::INTER_CUBIC);
    return floatMat;
}


std::vector<float> EnhanceWrapper::_infer(cv::Mat transDim){
    aclError ret = aclrtMemcpy(deviceInputBuffer, deviceInputBufferSize, transDim.data, deviceInputBufferSize, ACL_MEMCPY_HOST_TO_DEVICE);
    if(ret != ACL_SUCCESS){
        std::cerr << "acl copy data to host fail " << ret << std::endl;
        return {};
    }

    ret = aclmdlExecute(modelId, deviceInputDataset, deviceOutputDataset);
    if(ret != ACL_SUCCESS){
        std::cerr << "acl infer fail " << ret << std::endl;
        return {};
    }

    ret = aclrtMemcpy(hostOutBuffer, deviceOutputBufferSize, deviceOutputBuffer, deviceOutputBufferSize, ACL_MEMCPY_DEVICE_TO_HOST);
    if(ret != ACL_SUCCESS){
        std::cerr << "acl copy data to host fail " << ret << std::endl;
        return {};
    }

    float *outData = reinterpret_cast<float*>(hostOutBuffer);
    return std::vector<float>(outData, (outData + deviceOutputBufferSize));
}


std::pair<cv::Mat, cv::Size> EnhanceWrapper::preprocess(std::string imageSrc){
    return NetworkTool::ToolImagePreprocess(imageSrc);
}


EnhanceWrapper::~EnhanceWrapper(){
    aclrtFreeHost(hostOutBuffer);
    aclrtFree(deviceInputBuffer);
    aclmdlDestroyDataset(deviceInputDataset);
    aclrtFree(deviceOutputBuffer);
    aclmdlDestroyDataset(deviceOutputDataset);
    
    aclError ret = aclmdlDestroyDesc(modelDesc);
    if(ret != ACL_SUCCESS){
        std::cerr << "acl destroy desc fail， code is: " + std::to_string(ret) << std::endl;
    }

    ret = aclmdlUnload(modelId);
    if(ret != ACL_SUCCESS){
        std::cerr << "acl unload model fail， code is: " + std::to_string(ret) << std::endl;
    }

    if(stream != nullptr){
        ret = aclrtDestroyStream(stream);
        if(ret != ACL_SUCCESS){
            std::cerr << "acl destroy stream fail， code is: " + std::to_string(ret) << std::endl;
        }
    }
    
    if(context != nullptr){
        ret = aclrtDestroyContext(context);
        if(ret != ACL_SUCCESS){
            std::cerr << "acl destroy context fail， code is: " + std::to_string(ret) << std::endl;
        }
    }
    
    ret = aclrtResetDevice(0);
    if(ret != ACL_SUCCESS){
        std::cerr << "acl reset device fail， code is: " + std::to_string(ret) << std::endl;
    }
    
    ret = aclFinalize();
    if(ret != ACL_SUCCESS){
        std::cerr << "acl finalize fail， code is: " + std::to_string(ret) << std::endl;
    }
}

void EnhanceWrapper::initModelResource(){
    // device input node
    deviceInputDataset = aclmdlCreateDataset();
    deviceInputBufferSize = aclmdlGetInputSizeByIndex(modelDesc, 0);
    aclrtMalloc(&deviceInputBuffer, deviceInputBufferSize, ACL_MEM_MALLOC_HUGE_FIRST);
    aclDataBuffer *inputData = aclCreateDataBuffer(deviceInputBuffer, deviceInputBufferSize);
    aclError ret = aclmdlAddDatasetBuffer(deviceInputDataset, inputData);
    if(ret != ACL_SUCCESS){
        throw std::runtime_error("acl allocate input buffer fail， code is: " + std::to_string(ret));
    }

    // device output node
    deviceOutputDataset = aclmdlCreateDataset();
    deviceOutputBufferSize = aclmdlGetOutputSizeByIndex(modelDesc, 0);
    aclrtMalloc(&deviceOutputBuffer, deviceOutputBufferSize, ACL_MEM_MALLOC_HUGE_FIRST);
    aclDataBuffer *outputData = aclCreateDataBuffer(deviceOutputBuffer, deviceOutputBufferSize);
    ret = aclmdlAddDatasetBuffer(deviceOutputDataset, outputData);
    if(ret != ACL_SUCCESS){
        throw std::runtime_error("acl allocate output buffer fail， code is: " + std::to_string(ret));
    }

    // host output buffer
    aclrtMallocHost(&hostOutBuffer, deviceOutputBufferSize);
}


void EnhanceWrapper::initCannContext(std::string modelPath){
    aclError ret = aclInit(NULL);
    if(ret != ACL_SUCCESS){
        throw std::runtime_error("acl init failed， code is: " + std::to_string(ret));
    }

    ret = aclrtSetDevice(0);
    if(ret != ACL_SUCCESS){
        throw std::runtime_error("acl set device failed， code is: " + std::to_string(ret));
    }

    ret = aclrtCreateContext(&context, 0);
    if(ret != ACL_SUCCESS){
        throw std::runtime_error("acl create context failed， code is: " + std::to_string(ret));
    }


    ret = aclrtCreateStream(&stream);
    if(ret != ACL_SUCCESS){
        throw std::runtime_error("acl create stream failed， code is: " + std::to_string(ret));
    }

    ret = aclmdlLoadFromFile(modelPath.c_str(), &modelId);
    if(ret != ACL_SUCCESS){
        throw std::runtime_error("acl load model failed， code is: " + std::to_string(ret));
    }
    
    modelDesc = aclmdlCreateDesc();
    ret = aclmdlGetDesc(modelDesc, modelId);
    if(ret != ACL_SUCCESS){
        throw std::runtime_error("acl get model desc failed， code is: " + std::to_string(ret));
    }

    ret = aclrtGetRunMode(&runMode);
    if(ret != ACL_SUCCESS){
        throw std::runtime_error("acl get npu mode failed， code is: " + std::to_string(ret));
    }

}